{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "956fe0ad-d7cf-420d-8a60-8f54178f6f67",
   "metadata": {
    "tags": []
   },
   "source": [
    "# TransNormer-7B Lora微调"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b78e6035",
   "metadata": {},
   "source": [
    "## Step.1 环境加载 和 环境配置"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0a38ec9d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1）导入基础环境\n",
    "import pandas as pd\n",
    "import torch\n",
    "from datasets import Dataset\n",
    "\n",
    "# 2）导入大模型环境\n",
    "from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorForSeq2Seq, TrainingArguments, Trainer, GenerationConfig\n",
    "\n",
    "# 3）导入参数高效微调库\n",
    "from peft import LoraConfig, TaskType, get_peft_model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "906f7ffd-6b8b-48ec-a099-aa86a04b25de",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# 设置设备参数\n",
    "DEVICE = \"cuda\"  # 使用CUDA\n",
    "DEVICE_ID = \"0\"  # CUDA设备ID，如果未设置则为空\n",
    "CUDA_DEVICE = f\"{DEVICE}:{DEVICE_ID}\" if DEVICE_ID else DEVICE  # 组合CUDA设备信息"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "74f5962f",
   "metadata": {},
   "source": [
    "## Step.2 数据加载"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "8227a7a4-e44e-4dcf-a1d5-7b80135c3800",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The part of ds:\n",
      " {'instruction': ['小姐，别的秀女都在求中选，唯有咱们小姐想被撂牌子，菩萨一定记得真真儿的——', '这个温太医啊，也是古怪，谁不知太医不得皇命不能为皇族以外的人请脉诊病，他倒好，十天半月便往咱们府里跑。', '嬛妹妹，刚刚我去府上请脉，听甄伯母说你来这里进香了。'], 'input': ['', '', ''], 'output': ['嘘——都说许愿说破是不灵的。', '你们俩话太多了，我该和温太医要一剂药，好好治治你们。', '出来走走，也是散心。']}\n"
     ]
    }
   ],
   "source": [
    "# 1）将JSON文件转换为CSV文件\n",
    "df = pd.read_json('./huanhuan.json')\n",
    "ds = Dataset.from_pandas(df)\n",
    "\n",
    "print(\"The part of ds:\\n\",ds[:3])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "6739834a-8f08-43cc-8c49-8d531e3657f4",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "BaiChuanTokenizer(name_or_path='/root/autodl-tmp/OpenNLPLab/TransNormerLLM-7B/', vocab_size=64000, model_max_length=1000000000000000019884624838656, is_fast=False, padding_side='right', truncation_side='right', special_tokens={'bos_token': '<s>', 'eos_token': '</s>', 'unk_token': '<unk>'}, clean_up_tokenization_spaces=False),  added_tokens_decoder={\n",
       "\t0: AddedToken(\"<unk>\", rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),\n",
       "\t1: AddedToken(\"<s>\", rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),\n",
       "\t2: AddedToken(\"</s>\", rstrip=False, lstrip=False, single_word=False, normalized=True, special=True),\n",
       "}"
      ]
     },
     "execution_count": 4,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## 加载字符令牌\n",
    "tokenizer = AutoTokenizer.from_pretrained('/root/autodl-tmp/OpenNLPLab/TransNormerLLM-7B/', use_fast=False, trust_remote_code=True)\n",
    "tokenizer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "2bfefdee-6cdc-4c37-9e91-4f3d40bdc36d",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "tokenizer.pad_token = tokenizer.eos_token"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "7f5efe31-4358-4d2b-9d83-c38f81aac87b",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "## 定义数据令牌化函数\n",
    "def process_func(example):\n",
    "    MAX_LENGTH = 384    # Llama分词器会将一个中文字切分为多个token，因此需要放开一些最大长度，保证数据的完整性\n",
    "    input_ids, attention_mask, labels = [], [], []\n",
    "    instruction = tokenizer(f\"<|im_start|>system\\n现在你要扮演皇帝身边的女人--甄嬛<|im_end|>\\n<|im_start|>user\\n{example['instruction'] + example['input']}<|im_end|>\\n<|im_start|>assistant\\n\", add_special_tokens=False)  # add_special_tokens 不在开头加 special_tokens\n",
    "    response = tokenizer(f\"{example['output']}\", add_special_tokens=False)\n",
    "    input_ids = instruction[\"input_ids\"] + response[\"input_ids\"] + [tokenizer.pad_token_id]\n",
    "    attention_mask = instruction[\"attention_mask\"] + response[\"attention_mask\"] + [1]  # 因为eos token咱们也是要关注的所以 补充为1\n",
    "    labels = [-100] * len(instruction[\"input_ids\"]) + response[\"input_ids\"] + [tokenizer.pad_token_id]  \n",
    "    if len(input_ids) > MAX_LENGTH:  # 做一个截断\n",
    "        input_ids = input_ids[:MAX_LENGTH]\n",
    "        attention_mask = attention_mask[:MAX_LENGTH]\n",
    "        labels = labels[:MAX_LENGTH]\n",
    "    return {\n",
    "        \"input_ids\": input_ids,\n",
    "        \"attention_mask\": attention_mask,\n",
    "        \"labels\": labels\n",
    "    }"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "5e4b65d6-379b-47b0-b7bc-5967eeb50202",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "7aa1e8ccac024abb94161ab5cdef8931",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Map:   0%|          | 0/3729 [00:00<?, ? examples/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "Dataset({\n",
       "    features: ['input_ids', 'attention_mask', 'labels'],\n",
       "    num_rows: 3729\n",
       "})"
      ]
     },
     "execution_count": 7,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## 数据令牌化\n",
    "tokenized_id = ds.map(process_func, remove_columns=ds.column_names)\n",
    "tokenized_id"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "ac28c588-db5d-4153-93c4-d73340059d7d",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "' <|im_start|>system\\n现在你要扮演皇帝身边的女人--甄嬛<|im_end|>\\n<|im_start|>user\\n小姐，别的秀女都在求中选，唯有咱们小姐想被撂牌子，菩萨一定记得真真儿的——<|im_end|>\\n<|im_start|>assistant\\n 嘘——都说许愿说破是不灵的。</s>'"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "## 令牌化解码示例\n",
    "tokenizer.decode(tokenized_id[0]['input_ids'])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0f856ab3",
   "metadata": {},
   "source": [
    "## Step.3 模型加载&配置&Lora训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "39161659-95a4-4364-ac75-be94b59fda41",
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "2024-04-14 00:59:51 | INFO | accelerate.utils.modeling | We will use 90% of the memory on device 0 for storing the model, and 10% for the buffer to avoid OOM. You can set `max_memory` in to a higher value to use more memory (at your own risk).\n"
     ]
    },
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "01d516146b5c48bd9ed7b8665cc9d6f5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/3 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/root/miniconda3/lib/python3.8/site-packages/torch/_utils.py:776: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "TransnormerForCausalLM(\n",
       "  (model): TransnormerModel(\n",
       "    (embed_tokens): Embedding(64000, 4096, padding_idx=0)\n",
       "    (layers): ModuleList(\n",
       "      (0): TransnormerDecoderLayer(\n",
       "        (token_norm): SimpleRMSNorm()\n",
       "        (channel_norm): SimpleRMSNorm()\n",
       "        (token_mixer): NormLinearAttention(\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (norm): SimpleRMSNorm()\n",
       "          (lrpe): Lrpe((theta): Tensor((32, 1, 128), requires_grad=True))\n",
       "          (qkvu_proj): Linear(in_features=4096, out_features=16384, bias=False)\n",
       "        )\n",
       "        (channel_mixer): GLU(\n",
       "          (l1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (l2): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (l3): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "        )\n",
       "      )\n",
       "      (1-29): 29 x TransnormerDecoderLayer(\n",
       "        (token_norm): SimpleRMSNorm()\n",
       "        (channel_norm): SimpleRMSNorm()\n",
       "        (token_mixer): NormLinearAttention(\n",
       "          (out_proj): Linear(in_features=4096, out_features=4096, bias=False)\n",
       "          (norm): SimpleRMSNorm()\n",
       "          (qkvu_proj): Linear(in_features=4096, out_features=16384, bias=False)\n",
       "        )\n",
       "        (channel_mixer): GLU(\n",
       "          (l1): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (l2): Linear(in_features=4096, out_features=11008, bias=False)\n",
       "          (l3): Linear(in_features=11008, out_features=4096, bias=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "    (final_norm): SimpleRMSNorm()\n",
       "  )\n",
       "  (lm_head): Linear(in_features=4096, out_features=64000, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model = AutoModelForCausalLM.from_pretrained('/root/autodl-tmp//OpenNLPLab/TransNormerLLM-7B/',trust_remote_code=True, device_map=\"auto\",torch_dtype=torch.bfloat16)\n",
    "model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "2a236276-d7f3-49d7-a8e3-5ffbc246a3c8",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "config = LoraConfig(\n",
    "    task_type=TaskType.CAUSAL_LM, \n",
    "    target_modules=[\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"],\n",
    "    inference_mode=False, # 训练模式\n",
    "    r=8, # Lora 秩\n",
    "    lora_alpha=32, # Lora alaph，具体作用参见 Lora 原理\n",
    "    lora_dropout=0.1# Dropout 比例\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "f782f6fe-dcaf-4f0b-ba2e-6a814e137dbb",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "args = TrainingArguments(\n",
    "    output_dir=\"./output/DeepSeek\",\n",
    "    per_device_train_batch_size=4,\n",
    "    gradient_accumulation_steps=4,\n",
    "    logging_steps=10,\n",
    "    num_train_epochs=3,\n",
    "    save_steps=100,\n",
    "    learning_rate=1e-4,\n",
    "    save_on_each_node=True,\n",
    "    gradient_checkpointing=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9f741419-8341-4430-acf4-ce4f999e1ef7",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "trainer = Trainer(\n",
    "    model=model,\n",
    "    args=args,\n",
    "    train_dataset=tokenized_id,\n",
    "    data_collator=DataCollatorForSeq2Seq(tokenizer=tokenizer, padding=True),\n",
    ")\n",
    "trainer.train()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f284178a",
   "metadata": {},
   "source": [
    "## Step.4 Lora后的模型权重调用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "04c82e4c-ec84-4806-8898-ee3c48c25ede",
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoModelForCausalLM, AutoTokenizer\n",
    "import torch\n",
    "from peft import PeftModel\n",
    "\n",
    "mode_path = './OpenNLPLab/TransNormerLLM-7B/'\n",
    "lora_path = 'lora_path'\n",
    "\n",
    "# 加载tokenizer\n",
    "tokenizer = AutoTokenizer.from_pretrained(mode_path)\n",
    "\n",
    "# 加载模型\n",
    "model = AutoModelForCausalLM.from_pretrained(mode_path, device_map=\"auto\", torch_dtype=torch.bfloat16)\n",
    "\n",
    "# 加载lora权重\n",
    "model = PeftModel.from_pretrained(model, model_id=lora_path, config=config)\n",
    "\n",
    "prompt = \"你是谁？\"\n",
    "messages = [\n",
    "    {\"role\": \"system\", \"content\": \"现在你要扮演皇帝身边的女人--甄嬛\"},\n",
    "    {\"role\": \"user\", \"content\": prompt}\n",
    "]\n",
    "\n",
    "text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n",
    "\n",
    "model_inputs = tokenizer([text], return_tensors=\"pt\").to('cuda')\n",
    "\n",
    "generated_ids = model.generate(\n",
    "    model_inputs.input_ids,\n",
    "    max_new_tokens=512\n",
    ")\n",
    "generated_ids = [\n",
    "    output_ids[len(input_ids):] for input_ids, output_ids in zip(model_inputs.input_ids, generated_ids)\n",
    "]\n",
    "\n",
    "response = tokenizer.batch_decode(generated_ids, skip_special_tokens=True)[0]\n",
    "\n",
    "print(response)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.9 ('koopman_rl')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.9"
  },
  "vscode": {
   "interpreter": {
    "hash": "4e3e83c64c02740d29893152d09e5006b6c7a30ceaa232b939b9d06de7d1432d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
